import numpy as np

from simulation import NUMERIC_TYPE

class Integrator(object):
    def __init__(self, timestep=None):
        self.timestep = timestep
        # gravitational constant
        self.G = NUMERIC_TYPE(2.959122082855910E-04)

    def integrate(self, system, step):
        raise NotImplemented

    def accelerations(self, positions, masses):
        """ 
        Input: 
         positions is a dictionary of {planet : np.array(shape=(3)) }
         masses is a dictionary of { planet : mass } 
        Output:
         result is a dictionary of {planet : accelerations }
         energy is potential energy
        """
        acc = {}
        energy = NUMERIC_TYPE(0.0)
        for planet_i, mass in masses.iteritems():
            fp =  np.zeros(3, NUMERIC_TYPE) 
            for planet_j in masses.iterkeys():
                if planet_i != planet_j:
                    rv = positions[planet_i] - positions[planet_j]
                    r2 = np.dot(rv, rv)
                    r = r2 ** NUMERIC_TYPE(0.5)
                    d = self.G * masses[planet_j] / r
                    d3 = d / r2
                    fp -= np.dot(rv, d3)
                    energy -= d * mass # Energy is negative bc it is attractive
            acc[planet_i] = fp
        return acc, energy

    def forces(self, positions, masses):
        """
        Input: 
         positions is a dictionary of {planet : np.array(shape=(3)) }
         masses is a dictionary of { planet : mass } 
        Output:
         result is a dictionary of {planet : accelerations * mass }
         energy is potential energy
        """
        acc, pot_energy = self.accelerations(positions, masses)
        forces = {planet : acc[planet] * masses[planet] for planet \
                      in acc.iterkeys()}
        return forces, pot_energy

    def kinetic_energy(self, velocities, masses):
        result = NUMERIC_TYPE(0.0)
        for planet, mass in masses.iteritems():
            v = velocities[planet]
            v2 = np.dot(v, v)
            result += mass * v2
        return result*NUMERIC_TYPE(0.5)
